{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EKF\n",
    "使用EKF进行移动机器人的位姿预测，见相关[博客](https://blog.csdn.net/weixin_42301220/article/details/124605350?spm=1001.2014.3001.5501)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PyQt5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "%matplotlib qt5\n",
    "\n",
    "# # set up matplotlib\n",
    "# is_ipython = 'inline' in matplotlib.get_backend()\n",
    "# if is_ipython:\n",
    "#     from IPython import display\n",
    "\n",
    "plt.ion()\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as Rot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "设置过程噪声和观测噪声的协方差矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = np.matrix([[0.1, 0, 0, 0],  # variance of location on x-axis\n",
    "        [0, 0.1, 0, 0],  # variance of location on y-axis\n",
    "        [0, 0, np.deg2rad(1.0), 0],  # variance of yaw angle\n",
    "        [0, 0, 0, 1.0]])**2  # variance of velocity\n",
    "\n",
    "R = np.matrix([[1.0,0.0],[0.0,1.0]])**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "仿真参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Simulation parameter\n",
    "INPUT_NOISE = np.diag([1.0, np.deg2rad(30.0)]) ** 2\n",
    "GPS_NOISE = np.diag([0.5, 0.5]) ** 2\n",
    "\n",
    "DT = 0.1  # time tick [s]\n",
    "SIM_TIME = 50.0  # simulation time [s]\n",
    "\n",
    "show_animation = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "状态向量\n",
    "\n",
    "x_t = [x,y,yaw,v]\n",
    "\n",
    "控制输入为\n",
    "\n",
    "u = [v,w]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "控制输入表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def control_input():\n",
    "    v = 1.0  # 线速度，m/s\n",
    "    yaw_rate = 0.1  #角速度，rad/s\n",
    "    u = np.matrix([\n",
    "            [v],\n",
    "            [yaw_rate]\n",
    "            ])\n",
    "    return u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "观测模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observation_model(x):\n",
    "    H = np.matrix([[1.0,0,0,0],[0,1.0,0,0]])\n",
    "    z = H@x\n",
    "    return z\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "过程模型（状态方程）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def motion_model(x,u):\n",
    "    F = np.matrix([\n",
    "            [1.0,0,0,0],\n",
    "            [0,1.0,0,0],\n",
    "            [0,0,1.0,0],\n",
    "            [0,0,0,0]\n",
    "            ])\n",
    "    B = np.matrix([\n",
    "            [DT*np.cos(x[2,0]),0],\n",
    "            [DT*np.sin(x[2,0]),0],\n",
    "            [0,DT],\n",
    "            [1.0,0]\n",
    "            ])\n",
    "    x = F@x+B@u\n",
    "    return x\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "观测向量获取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observation(x_true, x_d,u):\n",
    "    x_true = motion_model(x_true,u)\n",
    "    # add noise to gps x-y\n",
    "    z = observation_model(x_true) + GPS_NOISE @ np.random.randn(2, 1)\n",
    "\n",
    "    # add noise to input\n",
    "    ud = u + INPUT_NOISE @ np.random.randn(2, 1)\n",
    "\n",
    "    xd = motion_model(x_d, ud)\n",
    "\n",
    "    return x_true, z, xd, ud\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "F的雅可比矩阵定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jacob_f(x,u):\n",
    "    \"\"\"\n",
    "    Jacobian of Motion Model\n",
    "\n",
    "    motion model\n",
    "    x_{t+1} = x_t+v*dt*cos(yaw)\n",
    "    y_{t+1} = y_t+v*dt*sin(yaw)\n",
    "    yaw_{t+1} = yaw_t+omega*dt\n",
    "    v_{t+1} = v{t}\n",
    "    so\n",
    "    dx/dyaw = -v*dt*sin(yaw)\n",
    "    dx/dv = dt*cos(yaw)\n",
    "    dy/dyaw = v*dt*cos(yaw)\n",
    "    dy/dv = dt*sin(yaw)\n",
    "    \"\"\"\n",
    "    yaw = x[2,0]\n",
    "    v = u[0,0]\n",
    "    J_F = np.matrix([\n",
    "            [1.0,0.0,-v*np.sin(yaw)*DT,np.cos(yaw)*DT],\n",
    "            [0.0,1.0,v*np.cos(yaw)*DT,np.sin(yaw)*DT],\n",
    "            [0,0,1,0],\n",
    "            [0,0,0,1]\n",
    "            ])\n",
    "    return J_F\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "H的雅可比矩阵定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jacob_h():\n",
    "    \"\"\"\n",
    "    Jacobian of Observation Model\n",
    "    \"\"\"\n",
    "    J_H = np.matrix([[1.0,0,0,0],[0,1.0,0,0]])\n",
    "    return J_H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ekf估计"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ekf_estimation(x_esti,P_esti,z,u):\n",
    "    \"\"\"两个阶段，预测和更新\n",
    "\n",
    "    Args:\n",
    "        x_esti (_type_): 估计的状态\n",
    "        P_esti (_type_): 估计的P矩阵（后验估计误差协方差矩阵）\n",
    "        z (_type_): 观测向量\n",
    "        u (_type_): 控制输入\n",
    "    \"\"\"\n",
    "    #  Predict\n",
    "    x_pred = motion_model(x_esti,u)\n",
    "    J_F = jacob_f(x_esti,u)\n",
    "    P_pred = J_F@P_esti@J_F.T+Q\n",
    "\n",
    "    # update\n",
    "    J_H = jacob_h()\n",
    "    z_pred = observation_model(x_pred)\n",
    "    y = z-z_pred\n",
    "    S = J_H@P_pred@J_H.T+R\n",
    "    K = P_pred@J_H.T@np.linalg.pinv(S)\n",
    "    x_esti = x_pred+K@y\n",
    "    P_esti = (np.eye(len(x_esti))-K@J_H)@P_pred\n",
    "    return x_esti,P_esti"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "画图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_covariance_ellipse(x_esti, P_esti):  # pragma: no cover\n",
    "    \"\"\"The blue line is true trajectory, the black line is dead reckoning trajectory,\n",
    "\n",
    "    the green point is positioning observation (ex. GPS), and the red line is estimated trajectory with EKF.\n",
    "\n",
    "    The red ellipse is estimated covariance ellipse with EKF.\n",
    "\n",
    "    Args:\n",
    "        x_esti (_type_): _description_\n",
    "        P_esti (_type_): _description_\n",
    "    \"\"\"\n",
    "    Pxy = P_esti[0:2, 0:2]\n",
    "    # 计算方形矩阵的特征值和特征向量\n",
    "    eigval, eigvec = np.linalg.eig(Pxy)\n",
    "\n",
    "    if eigval[0] >= eigval[1]:\n",
    "        bigind = 0\n",
    "        smallind = 1\n",
    "    else:\n",
    "        bigind = 1\n",
    "        smallind = 0\n",
    "\n",
    "    t = np.arange(0, 2 * math.pi + 0.1, 0.1)\n",
    "    a = math.sqrt(eigval[bigind])\n",
    "    b = math.sqrt(eigval[smallind])\n",
    "    x = [a * math.cos(it) for it in t]\n",
    "    y = [b * math.sin(it) for it in t]\n",
    "    angle = math.atan2(eigvec[1, bigind], eigvec[0, bigind])\n",
    "    rot = Rot.from_euler('z', angle).as_matrix()[0:2, 0:2]\n",
    "    fx = rot @ (np.array([x, y]))\n",
    "    px = np.array(fx[0, :] + x_esti[0, 0]).flatten()\n",
    "    py = np.array(fx[1, :] + x_esti[1, 0]).flatten()\n",
    "    plt.plot(px, py, \"--r\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "主函数入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def main():\n",
    "    nx = 4  # State Vector [x y yaw v]'\n",
    "    xEst = np.zeros((nx, 1))\n",
    "    xTrue = np.zeros((nx, 1))\n",
    "    PEst = np.eye(nx)\n",
    "    xDR = np.zeros((nx, 1))  # Dead reckoning\n",
    "\n",
    "\n",
    "    # history\n",
    "    hxEst = xEst\n",
    "    hxTrue = xTrue\n",
    "    hxDR = xTrue\n",
    "    hz = np.zeros((2, 1))\n",
    "\n",
    "    time = 0.0\n",
    "\n",
    "    while SIM_TIME >= time:\n",
    "        time += DT\n",
    "        u = control_input()\n",
    "\n",
    "        xTrue, z, xDR, ud = observation(xTrue, xDR, u)\n",
    "\n",
    "        xEst, PEst = ekf_estimation(xEst, PEst, z, ud)\n",
    "\n",
    "        # store data history\n",
    "        hxEst = np.hstack((hxEst, xEst))\n",
    "        hxDR = np.hstack((hxDR, xDR))\n",
    "        hxTrue = np.hstack((hxTrue, xTrue))\n",
    "        hz = np.hstack((hz, z))\n",
    "\n",
    "        if show_animation:\n",
    "            plt.cla()\n",
    "            # for stopping simulation with the esc key.\n",
    "            plt.gcf().canvas.mpl_connect('key_release_event', lambda event: [\n",
    "                exit(0) if event.key == 'escape' else None])\n",
    "            plt.plot(hz[0, :], hz[1, :], \".g\")\n",
    "            plt.plot(np.array(hxTrue[0, :]).flatten(),\n",
    "                     np.array(hxTrue[1, :]).flatten(), \"-b\")\n",
    "            plt.plot(np.array(hxDR[0, :]).flatten(),\n",
    "                     np.array(hxDR[1, :]).flatten(), \"-k\")\n",
    "            plt.plot(np.array(hxEst[0, :]).flatten(),\n",
    "                     np.array(hxEst[1, :]).flatten(), \"-r\")\n",
    "            plot_covariance_ellipse(xEst, PEst)\n",
    "            plt.axis(\"equal\")\n",
    "            plt.grid(True)\n",
    "            plt.pause(0.001)\n",
    "        plt.savefig('./result.png')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "main()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0c7484b3574347463e16b31029466871583b0d4e5c4ad861e8848f2d3746b4de"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('gobigger')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
